package com.moyz.adi.common.rag;

import com.moyz.adi.common.enums.ErrorEnum;
import com.moyz.adi.common.exception.BaseException;
import com.moyz.adi.common.interfaces.IRAGService;
import com.moyz.adi.common.util.InputAdaptor;
import com.moyz.adi.common.vo.InputAdaptorMsg;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.comparison.IsEqualTo;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.regex.Matcher;

import static com.moyz.adi.common.cosntant.AdiConstant.*;
import static com.moyz.adi.common.vo.InputAdaptorMsg.TOKEN_TOO_MUCH_QUESTION;
import static dev.langchain4j.model.openai.OpenAiModelName.GPT_3_5_TURBO;

@Slf4j
public class EmbeddingRAG implements IRAGService {

    private EmbeddingModel embeddingModel;

    private EmbeddingStore<TextSegment> embeddingStore;

    public EmbeddingRAG(EmbeddingStore<TextSegment> embeddingStore) {
        this.embeddingStore = embeddingStore;
    }

    public void init() {
        log.info("initEmbeddingModel");
        embeddingModel = new AllMiniLmL6V2EmbeddingModel();
    }

    /**
     * 对文档切块、向量化并存储到数据库
     *
     * @param document 知识库文档
     * @param overlap  重叠token数
     */
    @Override
    public void ingest(Document document, int overlap, ChatLanguageModel chatLanguageModel) {
        log.info("EmbeddingRAG ingest");
        DocumentSplitter documentSplitter = DocumentSplitters.recursive(RAG_MAX_SEGMENT_SIZE_IN_TOKENS, overlap, new OpenAiTokenizer(GPT_3_5_TURBO));
        EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(documentSplitter)
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        embeddingStoreIngestor.ingest(document);
    }

    /**
     * 创建一个基于嵌入存储的内容检索器。
     *
     * @param metadataCond           元数据条件，用于过滤要召回的文档。
     * @param maxResults             最大召回结果数量。
     * @param minScore               召回文档的最小匹配分数。如果提供的minScore小于等于0，则使用默认的RAG_MIN_SCORE。
     * @param breakIfSearchMissed    如果向量数据库中搜索不到数据，是否强行中断该搜索，不继续往下执行（即不继续请求LLM进行回答）。
     * @return 返回配置好的AdiEmbeddingStoreContentRetriever实例。
     */
    @Override
    public AdiEmbeddingStoreContentRetriever createRetriever(Map<String, String> metadataCond, int maxResults, double minScore, boolean breakIfSearchMissed) {
        // 初始化过滤器对象为null，准备构建复合过滤条件。
        Filter filter = null;
        // 遍历metadataCond中的每一项，构建AND连接的过滤条件。
        for (Map.Entry<String, String> entry : metadataCond.entrySet()) {
            String key = entry.getKey();
            String value = entry.getValue();
            // 如果filter为空，则初始化为第一个条件；否则将新条件与现有filter做AND组合。
            if (null == filter) {
                filter = new IsEqualTo(key, value);
            } else {
                filter = filter.and(new IsEqualTo(key, value));
            }
        }
        // 构建并返回AdiEmbeddingStoreContentRetriever实例，设置相关参数：
        return AdiEmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)          // 设置嵌入存储服务。
                .embeddingModel(embeddingModel)          // 设置嵌入模型。
                .maxResults(maxResults)                  // 设置最大召回结果数量。
                .minScore(minScore <= 0 ? RAG_MIN_SCORE : minScore) // 设置最小匹配分数，若minScore无效则使用默认值。
                .filter(filter)                          // 设置构建好的过滤条件。
                .breakIfSearchMissed(breakIfSearchMissed)// 设置搜索未命中时是否中断标志。
                .build();                                // 构建并返回检索器实例。
    }

    public static final String parsePromptTemplate(String question, String information) {
        return PROMPT_TEMPLATE.apply(Map.of("question", question, "information", Matcher.quoteReplacement(information))).text();
    }

    /**
     * 根据模型的contentWindow（即最大输入token数量）计算使用该模型最多可以召回的文档数量。
     * <br/>以分块时的最大文本段对应的token数量{RAG_MAX_SEGMENT_SIZE_IN_TOKENS}为计算因子。
     *
     * @param userQuestion   用户的问题文本。
     * @param maxInputTokens AI模型所能容纳的最大输入token数量。
     * @return 返回可以召回的最大文档数量。
     * @throws BaseException 如果maxInputTokens为0，则抛出参数错误异常。
     */
    public static int getRetrieveMaxResults(String userQuestion, int maxInputTokens) {

        // 检查AI模型的最大输入token数量是否有效。
        if (maxInputTokens == 0) {
            throw new BaseException(ErrorEnum.A_PARAMS_ERROR);
        }

        // 验证用户问题的有效性，并获取关于问题token数量的信息。
        InputAdaptorMsg inputAdaptorMsg = InputAdaptor.isQuestionValid(userQuestion, maxInputTokens);

        // 如果用户问题太长，以至于没有足够的token留给知识库召回的内容，则返回0表示不进行召回。
        if (inputAdaptorMsg.getTokenTooMuch() == TOKEN_TOO_MUCH_QUESTION) {
            log.warn("用户问题太长了，没有足够的token数量留给知识库召回的内容");
            return 0;
        } else {
            // 计算可用于召回文档的最大token长度。
            int maxRetrieveDocLength = maxInputTokens - inputAdaptorMsg.getUserQuestionTokenCount();

            // 如果可用的token长度超过最大召回数量乘以每个文档段的最大token数，
            // 则返回最大召回数量；否则，计算实际可以召回的文档数量。
            if (maxRetrieveDocLength > RAG_RETRIEVE_NUMBER_MAX * RAG_MAX_SEGMENT_SIZE_IN_TOKENS) {
                return RAG_RETRIEVE_NUMBER_MAX;
            } else {
                // 计算可以召回的文档数量，确保每个文档段不超过最大token数。
                return maxRetrieveDocLength / RAG_MAX_SEGMENT_SIZE_IN_TOKENS;
            }
        }
    }
}
